use std::{
    path::{Path, PathBuf},
    sync::Arc,
};

use bytemuck::{Pod, Zeroable, bytes_of};
use cudarc::driver::{
    CudaContext, CudaFunction, CudaModule, CudaStream, CudaView, CudaViewMut, DevicePtr,
    DriverError, LaunchConfig, PushKernelArg,
};

#[derive(Debug, Clone)]
pub struct Device {
    context: Arc<CudaContext>,
    default_stream: Arc<CudaStream>,
}

impl Device {
    pub fn new() -> Result<Self, DriverError> {
        let context = CudaContext::new(0)?;
        let default_stream = context.default_stream();
        Ok(Self {
            context,
            default_stream,
        })
    }

    pub fn load_shader_from_src<P: AsRef<Path>>(&self, path: P) -> Result<Shader, DriverError> {
        let ptx = cudarc::nvrtc::compile_ptx(std::fs::read_to_string(path).unwrap()).unwrap();
        let module = self.context.load_module(ptx)?;
        Ok(Shader { module })
    }

    pub fn load_default_shader(&self) -> Result<Shader, DriverError> {
        self.load_shader_from_src(
            PathBuf::from(file!())
                .parent()
                .unwrap()
                .join("cuda_shader.cu"),
        )
    }

    pub fn default_stream(&self) -> &Arc<CudaStream> {
        &self.default_stream
    }
}

#[derive(Debug, Clone)]
pub struct Shader {
    module: Arc<CudaModule>,
}

#[repr(C)]
#[derive(Clone, Copy, Zeroable, Pod)]
struct RasterizationKernelParam {
    modelview_mat: u64,
    normal_mat: u64,

    model_verts: u64,
    model_face_vrts: u64,

    model_norms: u64,
    model_face_norm: u64,

    model_tex: u64,
    model_face_tex: u64,

    face_num: u64,

    normal_map: u64,
    normal_map_width: u64,
    normal_map_height: u64,

    diffuse_map: u64,
    diffuse_map_width: u64,
    diffuse_map_height: u64,

    frame_buffer: u64,
    z_buffer: u64,
    fb_width: u64,
    fb_height: u64,
}

impl Shader {
    const MAT4_TRANSFORM_KERNEL_NAME: &'static str = "mat4_transform";
    const MAT4_TRANSFORM_INPLACE_KERNEL_NAME: &'static str = "mat4_transform_inplace";
    const RASTERIZATION_KERNEL_NAME: &'static str = "rasterization";
    const VERTEX_NDC_TRANSFORM_KERNEL_NAME: &'static str = "vertex_ndc_transform";

    fn load_function(&self, fn_name: &str) -> Result<CudaFunction, DriverError> {
        self.module.load_function(fn_name)
    }

    pub fn mat4_transform(
        &self,
        vertices: CudaView<f32>,
        mut transformed: CudaViewMut<f32>,
        mat: CudaView<f32>,
        stream: &Arc<CudaStream>,
    ) -> Result<(), DriverError> {
        assert!(vertices.len() % 4 == 0);
        assert!(vertices.len() == transformed.len());
        assert!(mat.len() == 4 * 4);

        let size = (vertices.len() / 4) as u32;

        let function = self.load_function(Self::MAT4_TRANSFORM_KERNEL_NAME)?;
        let mut builder = stream.launch_builder(&function);
        builder.arg(&vertices);
        builder.arg(&mut transformed);
        builder.arg(&mat);
        builder.arg(&size);
        unsafe { builder.launch(LaunchConfig::for_num_elems(size)) }?;

        Ok(())
    }

    pub fn mat4_transform_inplace(
        &self,
        mut vertices: CudaViewMut<f32>,
        mat: CudaView<f32>,
        stream: &Arc<CudaStream>,
    ) -> Result<(), DriverError> {
        assert!(vertices.len() % 4 == 0);
        assert!(mat.len() == 4 * 4);

        let size = (vertices.len() / 4) as u32;

        let function = self.load_function(Self::MAT4_TRANSFORM_INPLACE_KERNEL_NAME)?;
        let mut builder = stream.launch_builder(&function);
        builder.arg(&mut vertices);
        builder.arg(&mat);
        builder.arg(&size);
        unsafe { builder.launch(LaunchConfig::for_num_elems(size)) }?;

        Ok(())
    }

    pub fn vertex_ndc_transform(
        &self,
        mut vertices: CudaViewMut<f32>,
        stream: &Arc<CudaStream>,
    ) -> Result<(), DriverError> {
        assert!(vertices.len() % 4 == 0);

        let size = (vertices.len() / 4) as u32;

        let function = self.load_function(Self::VERTEX_NDC_TRANSFORM_KERNEL_NAME)?;
        let mut builder = stream.launch_builder(&function);
        builder.arg(&mut vertices);
        builder.arg(&size);
        unsafe { builder.launch(LaunchConfig::for_num_elems(size)) }?;

        Ok(())
    }

    pub fn rasterization(
        &self,
        model_view_mat: CudaView<f32>,
        normal_mat: CudaView<f32>,

        model_verts: CudaView<f32>,
        face_verts_indice: CudaView<u32>,

        norms: CudaView<f32>,
        face_norm_indice: CudaView<u32>,

        tex: CudaView<f32>,
        face_tex_indice: CudaView<u32>,

        normal_map: CudaView<u8>,
        normal_map_size: (usize, usize),

        diffuse_map: CudaView<u8>,
        diffuse_map_size: (usize, usize),

        framebuffer: CudaViewMut<u8>,
        z_buffer: CudaViewMut<f32>,
        fb_size: (usize, usize),
        stream: &Arc<CudaStream>,
    ) -> Result<(), DriverError> {
        assert!(model_verts.len() % 4 == 0);
        assert!(norms.len() == model_verts.len());
        assert!(face_verts_indice.len() % 3 == 0);
        assert!(face_norm_indice.len() % 3 == 0);
        assert!(tex.len() % 2 == 0);
        assert!(face_tex_indice.len() % 3 == 0);

        assert!(normal_map.len() == normal_map_size.0 * normal_map_size.1 * 4);

        assert!(framebuffer.len() == 4 * fb_size.0 * fb_size.1);
        assert!(z_buffer.len() == fb_size.0 * fb_size.1);

        let size = (face_verts_indice.len() / 3) as u64;
        let (fb_width, fb_height) = (fb_size.0 as u64, fb_size.1 as u64);

        let param = RasterizationKernelParam {
            modelview_mat: model_view_mat.device_ptr(stream).0,
            normal_mat: normal_mat.device_ptr(stream).0,

            model_verts: model_verts.device_ptr(stream).0,
            model_face_vrts: face_verts_indice.device_ptr(stream).0,

            model_norms: norms.device_ptr(stream).0,
            model_face_norm: face_norm_indice.device_ptr(stream).0,

            model_tex: tex.device_ptr(stream).0,
            model_face_tex: face_tex_indice.device_ptr(stream).0,

            face_num: size,

            normal_map: normal_map.device_ptr(stream).0,
            normal_map_width: normal_map_size.0 as u64,
            normal_map_height: normal_map_size.1 as u64,

            diffuse_map: diffuse_map.device_ptr(stream).0,
            diffuse_map_width: diffuse_map_size.0 as u64,
            diffuse_map_height: diffuse_map_size.1 as u64,

            frame_buffer: framebuffer.device_ptr(stream).0,
            z_buffer: z_buffer.device_ptr(stream).0,
            fb_width,
            fb_height,
        };
        let param_bytes = bytes_of(&param);
        let param_gpu = stream.memcpy_stod(param_bytes).unwrap();

        let function = self.load_function(Self::RASTERIZATION_KERNEL_NAME)?;
        let mut builder = stream.launch_builder(&function);
        builder.arg(&param_gpu);

        #[cfg(feature = "perf_test")]
        let t0 = std::time::Instant::now();
        #[cfg(feature = "perf_test")]
        {
            stream.synchronize().unwrap();
        }
        unsafe { builder.launch(LaunchConfig::for_num_elems(size as u32)) }?;
        #[cfg(feature = "perf_test")]
        {
            stream.synchronize().unwrap();
            let time_in_kernel = (std::time::Instant::now() - t0).as_secs_f32() * 1000.;
            println!(
                "execute kernel used time = {time_in_kernel} ms, expect fps = {}",
                1000. / time_in_kernel
            )
        }

        Ok(())
    }
}
